import numpy as np
from common import pprint


def calc_gray(rgb, imgType, precision=3):
    """
    rgb换算为灰度值值
    rgb: 元组
    precision: 保留小数位数，最多3位

    return: 灰度值
    """
    rgb = np.array(rgb)
    if imgType == 3:
        params = np.array([0.299, 0.587, 0.114])
    else:
        params = np.array([0.299, 0.587, 0.114, 0])
    gray = rgb.dot(params)
    if precision:
        return round(gray, precision)
    return int(gray)
    

def calc_image(image, imgType=3, dataType=None, isSort=False, isGray=False, precision=3):
    """
    图片像素按灰度、rgb重新排列
    image: 图片数据
    precision: 灰度精度
    dataType: 数据类型（根据draw_mode判断） 一维数组，二维数数组
    isGray: 是否生成灰度
    isSort: 是否进行排序

    return: 像素点的rgb数组
    """
    image_data = np.array(image)
    height, width = image_data.shape[:2]

    color_data = {}
    for row in range(width):
        pprint(f'图像像素点计算中{format(row/width*100, ".2f")}%。。。。', end='\r')
        for col in range(height):
            rgb = tuple(image_data[col, row])
            color_data[rgb] = 1 if rgb not in color_data else color_data[rgb] + 1
    pprint('图像计算完成')


    # 将字典转为数组
    color_data = np.array([[item[0], item[1]] for item in color_data.items()], dtype=object)
        
    if isGray:
        pprint('图像灰度序生成中。。。', end='\r')
        gray_data = np.array([calc_gray(item[0], imgType, precision) for item in color_data])  # 进行灰度排序
        image_data = np.insert(color_data, 1, values=gray_data, axis=1)
        pprint('图像灰度序列已完成')
    else:
        gray_data = np.array([1]*color_data.shape[0])  # 进行灰度排序
        image_data = np.insert(color_data, 1, values=gray_data, axis=1)

    if isSort:
        pprint('数据排序中。。。', end='\r')
        index = np.lexsort(image_data.T[:2,])
        image_data = image_data[index]
        pprint('数据已排序')

    pprint('生成图像数据。。。', end='\r')
    image_data = [[item[0]]*item[2] for item in image_data]
    image_data = np.array([item for arr in image_data for item in arr])
    pprint('图像数据已生成')

    if dataType == 'row':
        image_data = image_data.reshape(height, width, imgType)
    elif dataType == 'col':
        image_data = image_data.reshape(width, height, imgType)
        image_data = image_data.transpose(1, 0, 2)
    return image_data
